package main

import (
	"crypto/tls"
	"fmt"
	"io"
	"net/http"
	"os"
	"runtime"
	"strings"
	"sync"
	"time"

	"github.com/google/go-containerregistry/pkg/authn"
	"github.com/google/go-containerregistry/pkg/crane"
	"github.com/google/go-containerregistry/pkg/name"
	v1 "github.com/google/go-containerregistry/pkg/v1"
	"github.com/google/go-containerregistry/pkg/v1/remote"
	"github.com/google/go-containerregistry/pkg/v1/tarball"
	"github.com/sirupsen/logrus"
)

// Client 容器镜像下载客户端
type Client struct {
	remoteOptions []remote.Option
	craneOptions  []crane.Option
	insecure      bool // 是否使用HTTP协议
}

// SetAuth 设置认证信息
func (c *Client) SetAuth(username, password string) {
	if username != "" && password != "" {
		auth := &authn.Basic{
			Username: username,
			Password: password,
		}
		c.remoteOptions = append(c.remoteOptions, remote.WithAuth(auth))
		c.craneOptions = append(c.craneOptions, crane.WithAuth(auth))
	}
}

// SetTransport 设置传输配置
func (c *Client) SetTransport(transport *http.Transport) {
	c.remoteOptions = append(c.remoteOptions, remote.WithTransport(transport))
	c.craneOptions = append(c.craneOptions, crane.WithTransport(transport))
}

// SetInsecure 设置是否使用HTTP协议
func (c *Client) SetInsecure(insecure bool) {
	c.insecure = insecure
	if insecure {
		// 创建支持HTTP的传输配置
		transport := &http.Transport{
			TLSClientConfig: &tls.Config{InsecureSkipVerify: true},
		}
		c.remoteOptions = append(c.remoteOptions, remote.WithTransport(transport))
		c.craneOptions = append(c.craneOptions, crane.WithTransport(transport), crane.Insecure)
	}
}

// Pull 下载并解压容器镜像
func (c *Client) Pull(syncCount int, registry, imageName, tag, arch string, printInfo bool, onlyGetTag bool, username, password string) error {
	logrus.Debugf("开始处理镜像: %s:%s, 架构: %s", imageName, tag, arch)

	// 设置认证
	c.SetAuth(username, password)

	// 保存原始镜像名用于文件命名
	originalImageName := imageName

	// 构建完整的镜像引用
	var imageRef string
	if registry != "" && registry != "registry-1.docker.io" {
		logrus.Infof("自定义registry: %s", registry)

		// 检测并处理协议前缀
		if strings.HasPrefix(registry, "http://") {
			c.SetInsecure(true)
			registry = strings.TrimPrefix(registry, "http://")
		} else {
			registry = strings.TrimPrefix(registry, "https://")
		}
		imageRef = fmt.Sprintf("%s/%s:%s", registry, imageName, tag)
	} else {
		// Docker Hub
		if !strings.Contains(imageName, "/") {
			imageName = "library/" + imageName
		}
		imageRef = fmt.Sprintf("%s:%s", imageName, tag)
	}

	logrus.Infof("镜像引用: %s", imageRef)

	// 解析镜像引用，如果是不安全的registry，使用name.Insecure选项
	var ref name.Reference
	var err error
	if c.insecure {
		ref, err = name.ParseReference(imageRef, name.Insecure)
	} else {
		ref, err = name.ParseReference(imageRef)
	}
	if err != nil {
		return fmt.Errorf("解析镜像引用失败: %w", err)
	}

	// 如果只获取标签列表
	if onlyGetTag {
		return c.listTags(ref.Context())
	}

	// 获取镜像描述符
	desc, err := remote.Get(ref, c.remoteOptions...)
	if err != nil {
		// 检查是否为镜像不存在的错误
		errStr := err.Error()
		logrus.Errorf("获取镜像失败: %s", errStr)
		if strings.Contains(errStr, "authentication required") {
			return fmt.Errorf("需要登录认证 或 镜像不存在，请使用 -u 和 -p 参数提供用户名和密码尝试登录")
		}
		// 检查网络连接错误
		if strings.Contains(errStr, "EOF") || strings.Contains(errStr, "connection refused") || strings.Contains(errStr, "timeout") {
			return fmt.Errorf("网络连接失败，请检查网络连接或代理设置: %w", err)
		}
		return fmt.Errorf("获取镜像描述符失败: %w", err)
	}

	// 检查是否为多架构镜像
	if desc.MediaType.IsIndex() {
		return c.handleMultiArchImage(desc, arch, printInfo, originalImageName, tag)
	}

	// 如果不是多架构镜像，直接下载
	img, err := remote.Image(ref, c.remoteOptions...)
	if err != nil {
		// 检查是否为镜像不存在的错误
		errStr := err.Error()
		logrus.Errorf("获取镜像失败: %s", errStr)
		if strings.Contains(errStr, "authentication required") {
			return fmt.Errorf("需要登录认证 或 镜像不存在，请使用 -u 和 -p 参数提供用户名和密码尝试登录")
		}
		// 检查网络连接错误
		if strings.Contains(errStr, "EOF") || strings.Contains(errStr, "connection refused") || strings.Contains(errStr, "timeout") {
			return fmt.Errorf("网络连接失败，请检查网络连接或代理设置: %w", err)
		}
		return fmt.Errorf("获取镜像失败: %w", err)
	}

	// 获取镜像配置以确定架构
	configFile, err := img.ConfigFile()
	if err != nil {
		return fmt.Errorf("获取镜像配置失败: %w", err)
	}

	osArch := fmt.Sprintf("%s-%s", configFile.OS, configFile.Architecture)
	return c.downloadImage(img, originalImageName, tag, osArch, printInfo)
}

// listTags 列出镜像的所有标签
func (c *Client) listTags(repo name.Repository) error {
	tags, err := remote.List(repo, c.remoteOptions...)
	if err != nil {
		return fmt.Errorf("获取标签列表失败: %w", err)
	}

	fmt.Println("获取到的tag列表为:")
	fmt.Println(strings.Join(tags, "\n"))

	return nil
}

// handleMultiArchImage 处理多架构镜像
func (c *Client) handleMultiArchImage(desc *remote.Descriptor, arch string, printInfo bool, imageName, tag string) error {
	idx, err := desc.ImageIndex()
	if err != nil {
		return fmt.Errorf("获取镜像索引失败: %w", err)
	}

	manifest, err := idx.IndexManifest()
	if err != nil {
		return fmt.Errorf("获取索引清单失败: %w", err)
	}

	logrus.Infof("获得%d个架构信息:", len(manifest.Manifests))

	// 如果指定了架构，尝试匹配
	if arch != "" {
		for _, m := range manifest.Manifests {
			if m.Platform != nil && m.Platform.Architecture == arch {
				img, err := idx.Image(m.Digest)
				if err != nil {
					// 检查是否为镜像不存在的错误
					errStr := err.Error()
					if strings.Contains(errStr, "404") || strings.Contains(errStr, "not found") || strings.Contains(errStr, "MANIFEST_UNKNOWN") || strings.Contains(errStr, "NAME_UNKNOWN") {
						return fmt.Errorf("指定架构的镜像不存在: %s", arch)
					}
					if strings.Contains(errStr, "401") || strings.Contains(errStr, "403") || strings.Contains(errStr, "UNAUTHORIZED") || strings.Contains(errStr, "authentication required") {
						return fmt.Errorf("认证失败：该镜像需要登录认证，请使用 -u 和 -p 参数提供用户名和密码")
					}
					return fmt.Errorf("获取指定架构镜像失败: %w", err)
				}
				osArch := fmt.Sprintf("%s-%s", m.Platform.OS, m.Platform.Architecture)
				return c.downloadImage(img, imageName, tag, osArch, printInfo)
			}
		}
		return fmt.Errorf("未找到架构 %s 的镜像", arch)
	}

	// 过滤有效的架构清单
	var validManifests []v1.Descriptor
	for _, m := range manifest.Manifests {
		if m.Platform != nil && m.Platform.OS != "unknown" && m.Platform.Architecture != "unknown" {
			validManifests = append(validManifests, m)
		}
	}

	if len(validManifests) == 0 {
		return fmt.Errorf("没有可用的镜像架构")
	}

	// 显示可选架构
	fmt.Println("检测到多架构镜像，请选择要下载的架构:")
	for i, m := range validManifests {
		fmt.Printf("[%d] %s/%s\n", i+1, m.Platform.OS, m.Platform.Architecture)
	}

	// 查找本机架构的索引
	var defaultChoice int = -1
	for i, m := range validManifests {
		if m.Platform.OS == runtime.GOOS && m.Platform.Architecture == runtime.GOARCH {
			defaultChoice = i + 1
			break
		}
	}

	if defaultChoice > 0 {
		fmt.Printf("请输入选择的序号 (1-%d) 直接回车则自动适配本机架构 [%d]: ", len(validManifests), defaultChoice)
	} else {
		fmt.Printf("请输入选择的序号 (1-%d): ", len(validManifests))
	}

	// 读取用户输入
	var input string
	fmt.Scanln(&input)

	var choice int
	if input == "" && defaultChoice > 0 {
		// 用户直接回车，使用默认选择
		choice = defaultChoice
		fmt.Printf("使用本机架构: %s/%s\n", validManifests[choice-1].Platform.OS, validManifests[choice-1].Platform.Architecture)
	} else {
		// 解析用户输入
		_, err := fmt.Sscanf(input, "%d", &choice)
		if err != nil || choice < 1 || choice > len(validManifests) {
			return fmt.Errorf("无效的选择: %s，请输入 1-%d 之间的数字", input, len(validManifests))
		}
		fmt.Printf("选择架构: %s/%s\n", validManifests[choice-1].Platform.OS, validManifests[choice-1].Platform.Architecture)
	}

	// 下载选择的架构
	selectedManifest := validManifests[choice-1]
	img, err := idx.Image(selectedManifest.Digest)
	if err != nil {
		return fmt.Errorf("获取选择架构镜像失败: %w", err)
	}

	osArch := fmt.Sprintf("%s-%s", selectedManifest.Platform.OS, selectedManifest.Platform.Architecture)
	return c.downloadImage(img, imageName, tag, osArch, printInfo)
}

// downloadImage 下载镜像并生成Docker可导入格式
func (c *Client) downloadImage(img v1.Image, imageName, tag, osArch string, printInfo bool) error {
	// 获取镜像配置
	configFile, err := img.ConfigFile()
	if err != nil {
		return fmt.Errorf("获取镜像配置失败: %w", err)
	}

	if printInfo {
		c.printImageInfo(configFile)
		return nil
	}

	// 使用 crane 直接保存为 Docker 格式
	var outputFile string
	if imageName != "" {
		cleanImageName := strings.ReplaceAll(imageName, "/", "_")
		cleanImageName = strings.ReplaceAll(cleanImageName, ":", "_")
		// 添加标签和架构信息到文件名
		outputFile = fmt.Sprintf("%s-%s-%s.tar.gz", cleanImageName, tag, osArch)

	} else {
		outputFile = fmt.Sprintf("image_%s.tar.gz", time.Now().Format("20060102150405"))

	}

	logrus.Infof("保存为Docker格式: %s", outputFile)

	// 计算镜像总大小并显示进度
	return c.downloadImageWithProgress(img, imageName, tag, outputFile)
}

// downloadImageWithProgress 带进度显示的下载镜像
func (c *Client) downloadImageWithProgress(img v1.Image, imageName, tag, outputFile string) error {
	// 获取镜像清单以计算总大小
	manifest, err := img.Manifest()
	if err != nil {
		return fmt.Errorf("获取镜像清单失败: %w", err)
	}

	// 计算总大小
	var totalSize int64
	for _, layer := range manifest.Layers {
		totalSize += layer.Size
	}
	// 加上配置文件大小
	totalSize += manifest.Config.Size

	fmt.Printf("镜像总大小: %.2f MB\n", float64(totalSize)/(1024*1024))
	fmt.Println("正在下载镜像...")

	// 创建进度跟踪
	var downloadedSize int64
	var mu sync.Mutex

	// 启动进度显示协程
	done := make(chan bool)
	go func() {
		ticker := time.NewTicker(500 * time.Millisecond)
		defer ticker.Stop()
		for {
			select {
			case <-done:
				return
			case <-ticker.C:
				mu.Lock()
				current := downloadedSize
				mu.Unlock()
				if totalSize > 0 {
					progress := float64(current) / float64(totalSize) * 100
					if progress > 100 {
						progress = 100
					}
					fmt.Printf("\r下载进度: %.2f%% (%.2f MB / %.2f MB)",
						progress,
						float64(current)/(1024*1024),
						float64(totalSize)/(1024*1024))
				}
			}
		}
	}()

	// 创建自定义的 Writer 来跟踪写入的字节数
	file, err := os.Create(outputFile)
	if err != nil {
		return fmt.Errorf("创建输出文件失败: %w", err)
	}
	defer file.Close()

	progressWriter := &progressWriter{
		writer: file,
		onProgress: func(n int64) {
			mu.Lock()
			downloadedSize += n
			mu.Unlock()
		},
	}

	// 使用 tarball.Write 替代 crane.Save 以便跟踪进度
	ref, err := name.ParseReference(fmt.Sprintf("%s:%s", imageName, tag))
	if err != nil {
		return fmt.Errorf("解析镜像引用失败: %w", err)
	}

	tag_ref, ok := ref.(name.Tag)
	if !ok {
		return fmt.Errorf("无法转换为标签引用")
	}

	err = tarball.Write(tag_ref, img, progressWriter)
	if err != nil {
		return fmt.Errorf("写入镜像文件失败: %w", err)
	}

	// 停止进度显示
	close(done)

	// 确保显示100%进度
	fmt.Printf("\r下载进度: 100.00%% (%.2f MB / %.2f MB)",
		float64(downloadedSize)/(1024*1024),
		float64(downloadedSize)/(1024*1024))
	fmt.Printf("\n下载完成: %s\n", outputFile)
	return nil
}

// progressWriter 实现带进度跟踪的 Writer
type progressWriter struct {
	writer     io.Writer
	onProgress func(int64)
}

func (pw *progressWriter) Write(p []byte) (n int, err error) {
	n, err = pw.writer.Write(p)
	if err == nil && pw.onProgress != nil {
		pw.onProgress(int64(n))
	}
	return n, err
}

// printImageInfo 打印镜像信息
func (c *Client) printImageInfo(config *v1.ConfigFile) {
	fmt.Printf("镜像架构: %s/%s\n", config.OS, config.Architecture)
	fmt.Printf("创建时间: %s\n", config.Created.Format(time.RFC3339))
	fmt.Printf("工作目录: %s\n", config.Config.WorkingDir)
	fmt.Printf("入口点: %v\n", config.Config.Entrypoint)
	fmt.Printf("命令: %v\n", config.Config.Cmd)
	fmt.Printf("环境变量: %v\n", config.Config.Env)
	fmt.Printf("暴露端口: %v\n", config.Config.ExposedPorts)
	fmt.Printf("卷: %v\n", config.Config.Volumes)
	fmt.Printf("标签: %v\n", config.Config.Labels)
}

// SetLogLevel 设置日志级别
func SetLogLevel(lvl logrus.Level) {
	logrus.SetLevel(lvl)
}
